package com.oritsh.dicom.scu;

import com.oritsh.dicom.constant.CGetInformatiionModel;
import com.oritsh.dicom.exception.GetException;
import org.dcm4che.data.Attributes;
import org.dcm4che.data.Tag;
import org.dcm4che.data.UID;
import org.dcm4che.data.VR;
import org.dcm4che.io.DicomInputStream;
import org.dcm4che.net.*;
import org.dcm4che.net.pdu.AAssociateRQ;
import org.dcm4che.net.pdu.PresentationContext;
import org.dcm4che.net.pdu.RoleSelection;
import org.dcm4che.net.service.BasicCStoreSCP;
import org.dcm4che.net.service.DicomServiceException;
import org.dcm4che.net.service.DicomServiceRegistry;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

/**
 * Created with IntelliJ IDEA.
 * User: zarra
 * Date: 14-2-6
 * Time: 上午9:49
 * Copyright Shanghai Orient Rain Information Technology Co.,Ltd.
 */
public class GetSCU extends DicomSCU {
    CGetInformatiionModel model = CGetInformatiionModel.StudyRoot;

    Attributes keys = new Attributes();

    protected int priority;

    private int cancelAfter;

    private DicomStreamStore dicomStreamStore;

    List<String> tsuids = new ArrayList<>();


    private BasicCStoreSCP storageSCP = new BasicCStoreSCP("*") {

        @Override
        protected void store(Association as, PresentationContext pc, Attributes rq,
                             PDVInputStream data, Attributes rsp)
                throws IOException {
            if (dicomStreamStore == null)
                return;

            String iuid = rq.getString(Tag.AffectedSOPInstanceUID);
            String cuid = rq.getString(Tag.AffectedSOPClassUID);
            String tsuid = pc.getTransferSyntax();
            try {
                dicomStreamStore.store(as, as.createFileMetaInformation(iuid, cuid, tsuid), data);
            } catch (Exception e) {
                throw new DicomServiceException(Status.ProcessingFailure, e);
            }

        }
    };

    public DicomStreamStore getDicomStreamStore() {
        return dicomStreamStore;
    }

    public void setDicomStreamStore(DicomStreamStore dicomStreamStore) {
        this.dicomStreamStore = dicomStreamStore;
    }

    public GetSCU(CGetInformatiionModel model, String aeCalling) {
        super(aeCalling);
        this.model = model;


        tsuids.add(UID.JPEGLossless);
        tsuids.add(UID.JPEGLSLossless);
        tsuids.addAll(Arrays.asList(IVR_LE_FIRST));
        addLevel(null);
    }

    public GetSCU(String aeCalling) {
        this(CGetInformatiionModel.StudyRoot, aeCalling);
    }

    public GetSCU() {
        this(CGetInformatiionModel.StudyRoot, "GETSCU");
    }

    @Override
    protected PresentationContext createPresentationContext() {
        return new PresentationContext(1, model.cuid, tsuids.toArray(new String[0]));
    }

    @Override
    protected AAssociateRQ createAAssociateRQ(ApplicationEntity aeCalling, ApplicationEntity aeCalled) {

        String [] tsuids = this.tsuids.toArray(new String[0]);
        AAssociateRQ rq = super.createAAssociateRQ(aeCalling, aeCalled);
        addOfferedStorageSOPClass(rq, UID.ComputedRadiographyImageStorage,tsuids);
        addOfferedStorageSOPClass(rq, UID.DigitalXRayImageStorageForPresentation,tsuids);
        addOfferedStorageSOPClass(rq, UID.DigitalXRayImageStorageForProcessing,tsuids);
        addOfferedStorageSOPClass(rq, UID.DigitalMammographyXRayImageStorageForPresentation,tsuids);
        addOfferedStorageSOPClass(rq, UID.DigitalMammographyXRayImageStorageForProcessing,tsuids);
        addOfferedStorageSOPClass(rq, UID.CTImageStorage,tsuids);
        addOfferedStorageSOPClass(rq, UID.UltrasoundMultiFrameImageStorage,tsuids);
        addOfferedStorageSOPClass(rq, UID.MRImageStorage,tsuids);
        addOfferedStorageSOPClass(rq, UID.UltrasoundImageStorage,tsuids);
        addOfferedStorageSOPClass(rq, UID.SecondaryCaptureImageStorage,tsuids);
        addOfferedStorageSOPClass(rq, UID.GrayscaleSoftcopyPresentationStateStorageSOPClass,tsuids);
        addOfferedStorageSOPClass(rq, UID.XRayAngiographicImageStorage,tsuids);
        addOfferedStorageSOPClass(rq, UID.XRayRadiofluoroscopicImageStorage,tsuids);
        addOfferedStorageSOPClass(rq, UID.NuclearMedicineImageStorage,tsuids);
        addOfferedStorageSOPClass(rq, UID.VLPhotographicImageStorage,tsuids);
        addOfferedStorageSOPClass(rq, UID.VideoPhotographicImageStorage,tsuids);
        addOfferedStorageSOPClass(rq, UID.BasicTextSRStorage,tsuids);
        addOfferedStorageSOPClass(rq, UID.EnhancedSRStorage,tsuids);
        addOfferedStorageSOPClass(rq, UID.KeyObjectSelectionDocumentStorage,tsuids);
        addOfferedStorageSOPClass(rq, UID.XRayRadiationDoseSRStorage,tsuids);
        addOfferedStorageSOPClass(rq, UID.EncapsulatedPDFStorage,tsuids);
        addOfferedStorageSOPClass(rq, UID.PositronEmissionTomographyImageStorage,tsuids);

        return rq;
    }

    @Override
    protected Device createDevice(Connection localConn, ApplicationEntity local) {
        Device device = super.createDevice(localConn, local);
        device.setDimseRQHandler(createServiceRegistry());
        return device;
    }

    private DicomServiceRegistry createServiceRegistry() {
        DicomServiceRegistry serviceRegistry = new DicomServiceRegistry();
        serviceRegistry.addDicomService(storageSCP);
        return serviceRegistry;
    }

    private AAssociateRQ addOfferedStorageSOPClass(AAssociateRQ rq, String cuid,String ... tsuid) {

        if (!rq.containsPresentationContextFor(cuid))
            rq.addRoleSelection(new RoleSelection(cuid, false, true));

        for (String ts : tsuids)
            rq.addPresentationContext(new PresentationContext(
                2 * rq.getNumberOfPresentationContexts() + 1, cuid, ts));
        return rq;
    }

    private void addLevel(String s) {
        if (s == null)
            s = this.model.level;
        keys.setString(Tag.QueryRetrieveLevel, VR.CS, s);
    }

    public DimseRSP getAsync(Attributes attributes) throws GetException {
        FutureDimseRSP rsp = new FutureDimseRSP(as.nextMessageID());
        rsp.setAutoCancel(cancelAfter);
        getAsync(attributes, rsp);
        return rsp;
    }

    public void getAsync(Attributes attributes, DimseRSPHandler rspHandler) throws GetException {
        try {
            Attributes request = new Attributes();
            request.addAll(keys);
            request.addAll(attributes);
            as.cget(model.cuid, priority, request, null, rspHandler);
        } catch (IOException | InterruptedException e) {
            throw new GetException(e.getMessage());
        }
    }

    public DicomInputStream get(Attributes attributes) throws GetException {

        if (dicomStreamStore == null)
            return null;
        try {
            String sopUID = attributes.getString(Tag.SOPInstanceUID);
            DimseRSP rsp = getAsync(attributes);
            int lastStatus = -1;
            while (rsp.next()) {
                Attributes command = rsp.getCommand();
                lastStatus = command.getInt(Tag.Status, -1);

                if (Status.isPending(lastStatus))
                    continue;
            }
            if (lastStatus != Status.Success) {
                throw new GetException("status:" + lastStatus + " is not Success");
            }
            return dicomStreamStore.fetch(sopUID);
        } catch (IOException | InterruptedException e) {
            throw new GetException(e.getMessage());
        }
    }
}
